package org.apache.ocean;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.logging.Logger;

import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;

/**
 * 
 * @author jasonr
 */
// TODO: test MultiShardQueryRewrite with query that gets rewritten
public class MultiShardQueryRewrite {
	public static Logger log = Logger.getLogger(MultiShardQueryRewrite.class.getName());
	Query original;
	List<SearchIndexShard> searchIndexShards;
	ThreadPoolExecutor threadPool;
	TimeoutState timeoutState;

	public MultiShardQueryRewrite(Query original, List<SearchIndexShard> searchIndexShards, ThreadPoolExecutor threadPool, TimeoutState timeoutState) {
		this.original = original;
		this.searchIndexShards = searchIndexShards;
		this.threadPool = threadPool;
		this.timeoutState = timeoutState;
	}

	public Query rewrite() throws Exception {
		List<Callable<List<Query>>> shardRewriteTasks = new ArrayList<Callable<List<Query>>>(searchIndexShards.size());
		for (SearchIndexShard searchIndexShard : searchIndexShards) {
			ShardRewriteTask shardRewriteTask = new ShardRewriteTask(searchIndexShard);
			shardRewriteTasks.add(shardRewriteTask);
		}
		List<Query> queries = new ArrayList<Query>();
		List<Future<List<Query>>> futures = threadPool.invokeAll(shardRewriteTasks);
		for (Future<List<Query>> future : futures) {
			queries.addAll(future.get());
		}
		Query[] queryArray = queries.toArray(new Query[0]);
		if (queryArray == null || queryArray.length == 0)
			return null;
		return queryArray[0].combine(queryArray);
	}

	public class ShardRewriteTask implements Callable<List<Query>> {
		List<MultiIndexSearcher> indexSearchers;

		public ShardRewriteTask(SearchIndexShard searchIndexShard) throws Exception {
			indexSearchers = searchIndexShard.createMultiIndexSearchers();
		}

		public List<Query> call() throws Exception {
			List<Query> rewrittenQueries = new ArrayList<Query>(indexSearchers.size());
			for (IndexSearcher indexSearcher : indexSearchers) {
				Query rewritten = indexSearcher.rewrite(original);
				rewrittenQueries.add(rewritten);
				if (!timeoutState.isOK())
					throw new GeneralTimeoutException(timeoutState);
			}
			return rewrittenQueries;
		}
	}
}
